from sklearn import linear_model, preprocessing


class LinearModel(object):
    def __init__(self, model_name):
        self.model = linear_models[model_name]


linear_models = {
    'LinearRegression': linear_model.LinearRegression(),
    'Ridge': linear_model.Ridge(),
    'Lars': linear_model.Lars,
    'Lasso': linear_model.Lasso(),
    'LassoLars': linear_model.LassoLars(),
    'ElasticNet': linear_model.ElasticNet(),
    'OrthogonalMatchingPursuit': linear_model.OrthogonalMatchingPursuit(),
    'BayesianRegression': linear_model.BayesianRidge(),
    'AutomaticRelevanceDetermination': linear_model.ARDRegression(),
    'LogisticRegression': linear_model.LogisticRegression(),
    'TweedieRegressor': linear_model.TweedieRegressor(),
    'PoissonRegressor': linear_model.PoissonRegressor(),
    'StochasticGradientDescent': linear_model.SGDRegressor(),
    'Perceptron': linear_model.Perceptron(),
    'PAARegressor': linear_model.PassiveAggressiveRegressor(),
    'RANSAC': linear_model.RANSACRegressor(),
    'TheilSenRegressor': linear_model.TheilSenRegressor(),
    'HuberRegressor': linear_model.HuberRegressor(),
    'PolynomialFeatures': preprocessing.PolynomialFeatures(),
}

